#!/usr/bin/env python3
import tkinter as tk
from tkinter import ttk, scrolledtext
import socketio
import asyncio
import threading
from asyncio.subprocess import PIPE
import json
import random
import string
import platform

class WorkerGUI:
    def __init__(self, root):
        self.title = 'Worker: HK-' + ''.join(random.choices(string.ascii_letters, k=8))
        self.root = root
        self.root.title(self.title)
        self.root.resizable(False, False)  # 禁止窗口调整大小

        self.namespace = '/api/worker-ws'
        
        # 创建异步锁
        self.ui_lock = asyncio.Lock()
        
        # 初始化Socket.IO客户端
        self.sio = socketio.AsyncClient()
        self.setup_socket_events()
        
        # 创建UI
        self.create_widgets()
        
        # 异步事件循环
        self.loop = asyncio.new_event_loop()
        
    def setup_socket_events(self):
        @self.sio.event
        async def connect():
            print("Connected to server")
            await self.update_scrolledtext("Connected to server\n")
            
        @self.sio.event
        async def disconnect():
            print("Disconnected from server")
            await self.update_scrolledtext("Disconnected from server\n")
            
        @self.sio.on('message', namespace=self.namespace)
        async def handle_message(msg):
            await self.update_scrolledtext(f"Received message from server: {msg}\n")
            
        @self.sio.on('run_script', namespace=self.namespace)
        async def handle_run_script(data):
            print("Received run_script from server")
            await self.update_scrolledtext(f"Received run_script from server: {data}\n")
            await self.execute_script(data)
            
    async def execute_script(self, data):
        try:
            # 创建子进程执行脚本
            process = await asyncio.create_subprocess_exec(
                'python3', data['script'], data['task_id'],
                stdout=PIPE, stderr=PIPE
            )
            
            # 读取输出
            stdout, stderr = await process.communicate()
            
            if stdout:
                await self.update_scrolledtext(f"STDOUT:\n{stdout.decode()}\n")
            if stderr:
                await self.update_scrolledtext(f"STDERR:\n{stderr.decode()}\n")
                
            await self.update_scrolledtext(f"Script finished with return code: {process.returncode}\n")
        except Exception as e:
            await self.update_scrolledtext(f"Error executing script: {str(e)}\n")
    
    async def update_scrolledtext(self, text):
        async with self.ui_lock:
            self.output_text.config(state=tk.NORMAL)
            self.output_text.insert(tk.END, text)
            self.output_text.see(tk.END)
            self.output_text.config(state=tk.DISABLED)
    
    def create_widgets(self):
        # 第1行
        row1 = tk.Frame(self.root)
        row1.pack(fill=tk.X, padx=5, pady=5)
        
        tk.Label(row1, text="Server Register URL").pack(side=tk.LEFT)
        self.server_url_entry = tk.Entry(row1, width=40)
        self.server_url_entry.insert(0, 'http://127.0.0.1:7001')
        self.server_url_entry.pack(side=tk.LEFT, padx=5)
        
        self.connect_btn = tk.Button(row1, text="Connect", command=self.on_connect)
        self.connect_btn.pack(side=tk.LEFT)
        
        # 第2行
        row2 = tk.Frame(self.root)
        row2.pack(fill=tk.X, padx=5, pady=5)
        tk.Label(row2, text="Node capability configuration").pack(side=tk.LEFT)
        
        # 第3行
        row3 = tk.Frame(self.root)
        row3.pack(fill=tk.X, padx=5, pady=5)
        tk.Label(row3, text="Maximum number of available CPUs").pack(side=tk.LEFT)
        self.cpu_entry = tk.Entry(row3, width=5)
        self.cpu_entry.pack(side=tk.LEFT, padx=5)
        self.cpu_entry.insert(0, "8")
        
        # 第4行
        row4 = tk.Frame(self.root)
        row4.pack(fill=tk.X, padx=5, pady=5)
        tk.Label(row4, text="Maximum available memory capacity").pack(side=tk.LEFT)
        self.mem_entry = tk.Entry(row4, width=5)
        self.mem_entry.pack(side=tk.LEFT, padx=5)
        self.mem_entry.insert(0, "128")
        
        # 第5行
        row5 = tk.Frame(self.root)
        row5.pack(fill=tk.X, padx=5, pady=5)
        tk.Label(row5, text="GPU available").pack(side=tk.LEFT)
        self.gpu_var = tk.BooleanVar()
        self.gpu_check = tk.Checkbutton(row5, variable=self.gpu_var)
        self.gpu_check.pack(side=tk.LEFT)
        
        # 第6行
        row6 = tk.Frame(self.root)
        row6.pack(fill=tk.X, padx=5, pady=5)
        tk.Label(row6, text="Docker available").pack(side=tk.LEFT)
        self.docker_var = tk.BooleanVar()
        self.docker_check = tk.Checkbutton(row6, variable=self.docker_var)
        self.docker_check.pack(side=tk.LEFT)
        
        # 第7行 - 输出区域
        row7 = tk.Frame(self.root)
        row7.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)
        
        self.output_text = scrolledtext.ScrolledText(
            row7, wrap=tk.WORD, state=tk.DISABLED, height=15
        )
        self.output_text.pack(fill=tk.BOTH, expand=True)
    
    def on_connect(self):
        url = self.server_url_entry.get().strip()
        if not url:
            return
            
        # 获取节点能力配置
        capabilities = {
            "worker_id" : self.title,
            "os"        : platform.system(),
            "version"   : platform.version(),
            "cpus"      : int(self.cpu_entry.get()),
            "memory"    : int(self.mem_entry.get()),
            "gpu"       : self.gpu_var.get(),
            "docker"    : self.docker_var.get()
        }
        
        # 在后台线程中运行异步连接
        def connect_async():
            async def connect():
                try:
                    await self.sio.connect(url, namespaces=self.namespace)
                    await self.sio.emit('register', data=capabilities, namespace=self.namespace)
                    await self.update_scrolledtext(f"Connected to {url}\n")
                except Exception as e:
                    asyncio.run_coroutine_threadsafe(
                        self.update_scrolledtext(f"Connection error: {str(e)}\n"),
                        self.loop
                    )
            
            asyncio.run_coroutine_threadsafe(connect(), self.loop)
        
        threading.Thread(target=connect_async, daemon=True).start()
    
    def run(self):
        # 启动异步事件循环
        def start_loop():
            asyncio.set_event_loop(self.loop)
            self.loop.run_forever()
        
        threading.Thread(target=start_loop, daemon=True).start()
        self.root.mainloop()

if __name__ == "__main__":
    root = tk.Tk()
    app = WorkerGUI(root)
    app.run()


# vim: set nu expandtab softtabstop=4 tabstop=4 tw=80:
